/***************************************************************************
 *   Copyright (C) 2008 by Eduardo Gurgel Pinho                            *
 *   edgurgel@gmail.com                                                    *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 3 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
***************************************************************************/

#include "kd_tree.h"

#include <vector>
#include <stack>
#include <cmath>
#include <queue>
#include <limits>

namespace Sirat {

namespace Photon_Mapping {

Kd_Node::Kd_Node(const Math::Vector& _point,unsigned short _dimension) : left(0), right(0), point(_point), dimension(_dimension) , photon(0) { }

Kd_Node::~Kd_Node(void) {
	if(photon != 0) delete photon;
	if(this->right != 0) delete right;
	if(this->left != 0) delete left;
}

void Kd_Node::add_left(Kd_Node* _left) { 
	left = _left; 
}

void Kd_Node::add_right(Kd_Node* _right) { 
	right = _right; 
}

void Kd_Node::add_photon(Photon* _photon) { 
	photon = _photon; 
}

Kd_Node* Kd_Node::get_left(void) const { 
	return left; 
}

Kd_Node* Kd_Node::get_right(void) const { 
	return right; 
}

Photon* Kd_Node::get_photon(void) const { 
	return photon; 
}

int Kd_Node::get_height(void) const {
	if(this == 0) {
		return -1;
	}
	int hl = this->right->get_height();
	int hr = this->left->get_height();
	int h = 1 + std::max(hl,hr);
	return h;
}

Kd_Tree::Kd_Tree(void) : root(0) { }

Kd_Tree::~Kd_Tree(void) {
	if(root != 0) delete root;
}

Photon* Kd_Tree::nearest_neighbour(const Math::Vector& _point) const {
	
	double root_distance = root->point.distance(_point);
	double minimal_distance = root_distance;
	Photon* root_photon = root->photon;
	Photon* nearest_photon = root_photon;
	
	std::stack< tuple<unsigned short,Kd_Node*,unsigned short> > s;
	/** Flags 
	 *	0 = Go ahead
	 *	1 = Look at left brother
	 *	2 = Look at right brother
	 */
	s.push(make_tuple(1,root,0));
	while(!(s.empty())) {
		tuple<unsigned short,Kd_Node*,unsigned short> current = s.top();
		s.pop();
		
		/* Tuple values */
		const unsigned short dimension = current.get<0>();
		Kd_Node* kd_node = current.get<1>();
		unsigned short flag = current.get<2>();
		const unsigned short next_dimension = (dimension != 3)?dimension + 1:1;
		
		const double value = _point[next_dimension];
		const double current_value = kd_node->point[next_dimension];
		
		switch(flag) {
			case 0: {
				Kd_Node* next_node = 0;
				Kd_Node* next_other_node = 0;
				unsigned short next_flag = 0;
				if(value <= current_value) {
					next_node = kd_node->left;
					next_other_node = kd_node->right;
					next_flag = 2;
				}
				else {
					next_node = kd_node->right;
					next_other_node = kd_node->left;
					next_flag = 1;
				}
				if(next_node != 0) {
					double next_distance = _point.distance(next_node->point);
					if(next_distance <= minimal_distance) {
						minimal_distance = next_distance;
						nearest_photon = next_node->photon;
					}
					s.push(make_tuple(dimension,kd_node,next_flag));
					s.push(make_tuple(next_dimension,next_node,0));
				}
				else if(next_other_node != 0) {
					double next_other_distance = _point.distance(next_other_node->point);
					if(next_other_distance <= minimal_distance) {
						minimal_distance = next_other_distance;
						nearest_photon = next_other_node->photon;
					}
					s.push(make_tuple(dimension,kd_node,next_flag));
					s.push(make_tuple(next_dimension,next_other_node,0));
				}
				break;
			}
			/* Left brother! */
			case 1: {
				if(kd_node->left != 0) {
					if(fabs(current_value - value) < minimal_distance) {
						double test_distance = _point.distance(kd_node->left->point);
						if(test_distance < minimal_distance) {
							minimal_distance = test_distance;
							nearest_photon = kd_node->left->photon;
						}
						s.push(make_tuple(next_dimension,kd_node->left,0));
					}
				}
				break;
			}
			/* Right brother */
			case 2: {
				if(kd_node->right != 0) {
					if(fabs(current_value - value) < minimal_distance) {
						double test_distance = _point.distance(kd_node->right->point);
						if(test_distance < minimal_distance) {
							minimal_distance = test_distance;
							nearest_photon = kd_node->right->photon;
						}
						s.push(make_tuple(next_dimension,kd_node->right,0));
					}
				}
				break;
			}	
		}
	}
	return nearest_photon;
}

std::vector<tuple<Photon*,double> > Kd_Tree::k_nearest_neighbour(const Math::Vector& _point,const unsigned short _k,const double _max_radius) const {
	std::numeric_limits<double> double_limit;
	double minimal_distance = double_limit.infinity();
	
	/**
	 * Limiting range.
	 */
	double max_radius = double_limit.infinity();
	if(_max_radius > 0) max_radius = _max_radius;
	
	double root_distance = root->point.distance(_point);
	Photon* root_photon = root->photon;

	
	std::vector<tuple<Photon*,double> > nearest;
	nearest.push_back(make_tuple(root_photon,root_distance));
	
	std::stack< tuple<unsigned short,Kd_Node*,unsigned short> > s;
	
	/** Flags 
	 * 0 = Go ahead
	 * 1 = Look at left brother
	 * 2 = Look at right brother
	 */
	
	s.push(make_tuple(1,root,0));
	while(!(s.empty())) {
		tuple<unsigned short,Kd_Node*,unsigned short> current = s.top();
		s.pop();
		
		/* Tuple values */
		const unsigned short dimension = current.get<0>();
		Kd_Node* kd_node = current.get<1>();
		unsigned short flag = current.get<2>();
		const unsigned short next_dimension = (dimension != 3)?dimension + 1:1;
		
		const double value = _point[next_dimension];
		const double current_value = kd_node->point[next_dimension];
	
		switch(flag) {
			case 0: {
				Kd_Node* next_node = 0;
				Kd_Node* next_other_node = 0;
				unsigned short next_flag = 0;
				if(value <= current_value) {
					next_node = kd_node->left;
					next_other_node = kd_node->right;
					next_flag = 2;
				}
				else {
					next_node = kd_node->right;
					next_other_node = kd_node->left;
					next_flag = 1;
				}
				if(next_node != 0) {
					double next_distance = _point.distance(next_node->point);
					Photon* next_photon = next_node->photon;
					if(next_distance <= minimal_distance && next_distance <= max_radius) {
						insert_knn_list(nearest,next_photon,next_distance,minimal_distance,_k);
					}
					s.push(make_tuple(dimension,kd_node,next_flag));
					s.push(make_tuple(next_dimension,next_node,0));
				}
				else if(next_other_node != 0) {
					double next_distance = _point.distance(next_other_node->point);
					Photon* next_photon = next_other_node->photon;
					if(next_distance <= minimal_distance && next_distance <= max_radius) insert_knn_list(nearest,next_photon,next_distance,minimal_distance,_k);
					s.push(make_tuple(dimension,kd_node,next_flag));
					s.push(make_tuple(next_dimension,next_other_node,0));
				}
				break;
			}
			/**
			 * Left brother! 
			 */
			case 1: {
				if(kd_node->left != 0) {
					if(fabs(current_value - value) < minimal_distance) {
						double next_distance = _point.distance(kd_node->left->point);
						Photon* next_photon = kd_node->left->photon;
						if(next_distance <= minimal_distance  && next_distance <= max_radius) insert_knn_list(nearest,next_photon,next_distance,minimal_distance,_k);
						s.push(make_tuple(next_dimension,kd_node->left,0));
					}
				}
				break;
			}
			/**
			 * Right brother 
			 */
			case 2: {
				if(kd_node->right != 0) {
					if(fabs(current_value - value) < minimal_distance) {
						double next_distance = _point.distance(kd_node->right->point);
						Photon* next_photon = kd_node->right->photon;
						if(next_distance <= minimal_distance && next_distance <= max_radius) insert_knn_list(nearest,next_photon,next_distance,minimal_distance,_k);
						s.push(make_tuple(next_dimension,kd_node->right,0));
					}
				}
				break;
			}	
		}
	}
	return nearest;
}

/**
 * This function tries to insert a possible nearest neighbour at this list 
 */
void insert_knn_list(std::vector<tuple<Photon*,double> >& nearest,Photon* next_photon,const double& next_distance,double& minimal_distance,const unsigned short _k) {
	if(nearest.size() == _k) {
		bool exist = false;
		/**
		 * Checking if this photon isn't at "nearests" .
		 */
		for(unsigned int i = 0; i < nearest.size(); i++)
			if(nearest[i].get<0>() == next_photon) { exist = true; break; } 
		if(exist == false) 
			for(std::vector<tuple<Photon*,double> >::iterator i = nearest.begin() ; i != nearest.end() ; i++ ) {
			if( next_distance < (*i).get<1>() ) {
				nearest.insert(i,make_tuple(next_photon,next_distance));
				nearest.erase(nearest.end()-1);
				minimal_distance = (*(nearest.end()-1)).get<1>();
				break;
			}
			}
	} else {
		bool insert = false;
		bool exist = false;
		/**
		 * Checking if this photon isnt at 'nearests' 
		 */
		for(unsigned int i = 0; i < nearest.size(); i++)
			if(nearest[i].get<0>() == next_photon) { exist = true; break; } 
		if(exist == false) {
			for(std::vector<tuple<Photon*,double> >::iterator i = nearest.begin() ; i != nearest.end() ; i++ ) {
				if( next_distance < (*i).get<1>() ) {
					nearest.insert(i,make_tuple(next_photon,next_distance));
					insert = true;
					break;
				}
			}
			if(insert == false) nearest.insert(nearest.end(),make_tuple(next_photon,next_distance));
		}
	}
}

int Kd_Tree::get_height(void) const {
	return 1 + root->get_height();
}

void Kd_Tree::level_order(void) const {
	std::queue<Kd_Node*> q;
	q.push(this->root);
	
	while(!(q.empty())) {
		Kd_Node* current = q.front();
		q.pop();
		if(current != 0) {
			current->point.show();
			q.push(current->left);
			q.push(current->right);
		}
	}
}

}

}
